#include "CameraCalibration.h"
#include <iostream>

CameraCalibration::CameraCalibration(const std::string& calibrationImagePath, const cv::Size& chessboardSize, float squareSize)
    : calibrationImagePath(calibrationImagePath), chessboardSize(chessboardSize), squareSize(squareSize) {
}

CameraCalibration::~CameraCalibration() {}

bool CameraCalibration::performCalibration() {
    std::vector<std::vector<cv::Point3f>> objectPoints;
    std::vector<std::vector<cv::Point2f>> imagePoints;
    std::vector<cv::String> filenames;

    // 获取标定图像文件名
    cv::glob(calibrationImagePath + "/*.jpg", filenames);
    if (filenames.empty()) {
        std::cerr << "No calibration images found." << std::endl;
        return false;
    }

    // 准备标定板角点的三维坐标
    std::vector<cv::Point3f> obj;
    for (int i = 0; i < chessboardSize.height; ++i) {
        for (int j = 0; j < chessboardSize.width; ++j) {
            obj.emplace_back(j * squareSize, i * squareSize, 0);
        }
    }

    // 遍历所有标定图像
    for (const auto& filename : filenames) {
        cv::Mat image = cv::imread(filename, cv::IMREAD_GRAYSCALE);
        if (image.empty()) {
            std::cerr << "Failed to read image: " << filename << std::endl;
            continue;
        }

        std::vector<cv::Point2f> corners;
        bool found = cv::findChessboardCorners(image, chessboardSize, corners);

        if (found) {
            // 亚像素级角点检测
            cv::cornerSubPix(image, corners, cv::Size(11, 11), cv::Size(-1, -1),
                cv::TermCriteria(cv::TermCriteria::EPS + cv::TermCriteria::MAX_ITER, 30, 0.1));
            objectPoints.push_back(obj);
            imagePoints.push_back(corners);

            // 绘制并显示角点
            cv::drawChessboardCorners(image, chessboardSize, corners, found);
            cv::imshow("Corners", image);
            cv::waitKey(500);
        }
    }

    cv::destroyAllWindows();

    if (objectPoints.empty()) {
        std::cerr << "No valid chessboard corners found in any of the images." << std::endl;
        return false;
    }

    // 进行相机标定
    std::vector<cv::Mat> rvecs, tvecs;
    reprojectionError = cv::calibrateCamera(objectPoints, imagePoints, imagePoints[0].size(),
        cameraMatrix, distortionCoefficients, rvecs, tvecs);

    std::cout << "Reprojection error: " << reprojectionError << std::endl;
    std::cout << "Camera matrix:" << std::endl << cameraMatrix << std::endl;
    std::cout << "Distortion coefficients:" << std::endl << distortionCoefficients << std::endl;

    return true;
}

cv::Mat CameraCalibration::getCameraMatrix() const {
    return cameraMatrix;
}

cv::Mat CameraCalibration::getDistortionCoefficients() const {
    return distortionCoefficients;
}

double CameraCalibration::getReprojectionError() const {
    return reprojectionError;
}